#!/bin/bash

# Copyright (c) 2024, Kuaishou Technology. All rights reserved.

set -euo pipefail

source ./llama-34b

export SEQ_LENGTH=4096

export HOSTFILE=$1
NUM_NODE=$2
export MASTER_ADDR=$(cat $HOSTFILE | head -n 1 | sed -s 's/slots=8//g')
export NUM_GPUS=$((NUM_NODE*8))
export GLOBAL_BATCH_SIZE=$((NUM_NODE*8))

export TP=2
export CP=1
export PP=4
export PP_l=1
export CKPT=full
export OFFLOAD_ALPHA=0

./pretrain_llama.sh
